import torch
from PIL import Image
from torchvision import transforms

from model import UnetGenerator
from utils import denormalize, GuidedFilter

infer_transform = transforms.Compose([
    # transforms.Resize((256, 256)),# inference  no need
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

'''
生成图片
'''


class Inference(object):
    def __init__(self):
        self.generator = UnetGenerator()
        model_state_dict = torch.load("./weights/generator/best.pth")['model_state_dict']
        new_state_dict = {}
        self.guided_filter = GuidedFilter()
        for k, v in model_state_dict.items():
            new_state_dict[k[7:]] = v

        self.generator.load_state_dict(new_state_dict)
        self.generator.eval()
        self.yita=0.5

    def inference(self, img_path):
        with torch.no_grad():
            img = Image.open(img_path)
            height = img.height
            width = img.width
            if not height % 2 == 0:
                height = height + 2
            if not width % 2 == 0:
                width = width + 2

            img = img.resize((width, height))

            img = infer_transform(img).unsqueeze(0)
            generator_img = self.generator(img)
            generator_guide_img = self.guided_filter( img,generator_img ,r=1)
            result=self.yita*generator_guide_img+(1-self.yita)*generator_img
            self.get_image(result, img)

    def get_image(self, image, photo):
        output = (denormalize(image.permute((0, 2, 3, 1)).
                              detach().
                              to('cpu').
                              numpy()) * 255).astype('uint8')
        output = output[0]

        photo = (denormalize(photo.permute((0, 2, 3, 1)).
                             detach().
                             to('cpu').
                             numpy()) * 255).astype('uint8')

        photo = photo[0]

        width = photo.shape[1]
        height = photo.shape[0]

        output = Image.fromarray(output).convert('RGB')
        photo = Image.fromarray(photo).convert("RGB")

        target = Image.new('RGB', (width + width, height), (255, 255, 255))

        target.paste(output, (0, 0, width, height))
        target.paste(photo, (width, 0, 2 * width, height))

        target.show()
        target.save("save.jpg")


model = Inference()
model.inference(r"E:\Datasets2\person_detect\person_from_武装分子\jpg\1efab6c2-cce5-11ea-bcf2-b2fc36340453.jpg")
